from pathml.ml.hovernet import compute_hv_map
import torch
import os
import glob
from PIL import Image
import cv2
import numpy as np

class MICCAIDataSet(torch.utils.data.Dataset):
    

    def __init__(self, file_path: str, mode='train', transform=None):
        """
            file_path (str): 数据集路径
            mode (str): train, val, test
            transform: 数据增强
        """
        if mode == 'train':
            self.img_path = glob.glob(os.path.join(file_path, 'TrainDataset/images512/*.jpg'))
            self.mask_path = glob.glob(os.path.join(file_path, 'TrainDataset/masks512/*.jpg'))
        elif mode == 'val':
            self.img_path = glob.glob(os.path.join(file_path, 'ValidationDataset/images512/*.jpg'))
            self.mask_path = glob.glob(os.path.join(file_path, 'ValidationDataset/masks512/*.jpg'))
        else:
            self.img_path = glob.glob(os.path.join(file_path, 'TestDataset/images512/*.jpg'))
            self.mask_path = glob.glob(os.path.join(file_path, 'TestDataset/masks512/*.jpg'))
        self.transform = transform
        self.mode = mode

    def __len__(self):
        return len(self.img_path)

    def __getitem__(self, idx):
        img = self.img_path[idx]
        mask = self.mask_path[idx]
        img = cv2.imread(img)
        mask = cv2.imread(mask)
        # print(mask[mask > 0])
        # Resize 图像和掩码
        img = cv2.resize(img, (512, 512))
        mask = cv2.resize(mask, (512, 512))
        # print(img.shape, mask.shape)
        if self.transform is not None:
            data = self.transform(image=img, mask=mask)
            img = data['image']
            mask = cv2.cvtColor(data['mask'], cv2.COLOR_BGR2GRAY)
        else:
            mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
            # mask[mask > 0] = 255
        mask = mask // 255
        _, mask_ = cv2.connectedComponents(mask)
        hv_map = compute_hv_map(mask_)
        mask = np.where(mask == 0, 1, 0)
        # r = self.img_path[idx].split('/')[-1]
        # cv2.imwrite(f'../data/MICCAI2018MoNuSeg/res/{r}111.png', mask)
        img = img.transpose([2, 0, 1])     
        img, mask, hv_map = torch.from_numpy(img), torch.from_numpy(mask).unsqueeze(0), torch.from_numpy(hv_map)
        if self.mode != 'test':
            return img, mask, hv_map
        return img, mask, self.img_path[idx].split('/')[-1]
        